﻿/*
 * H265PropParser.cpp
 *
 *  Created on: 2016年4月7日
 *      Author: terry
 */

#include "H265PropParser.h"

namespace av
{

const uint8_t H265PropParser::s_startCode[START_CODE_LENGTH] = { 0, 0, 0, 1};


H265PropParser::H265PropParser()
{
}

H265PropParser::~H265PropParser()
{
}

const uint8_t* H265PropParser::getStartCode()
{
    return s_startCode;

}

bool H265PropParser::startWithCode(const uint8_t* data, size_t length)
{
    if (length < 4)
    {
        return false;
    }

    return (data[0] == 0 && data[1] == 0 && data[2] == 0 && data[3] == 1);
}

size_t H265PropParser::findStartCode(const uint8_t* data, size_t length, size_t start)
{
    for (size_t i = start; i < length - 4; ++ i)
    {
        if (data[i] == 0)
        {
            if (startWithCode(&data[i], length - i))
            {
                return i;
            }
        }
    }
    return -1;
}

#define RB16(x) ((((uint8_t*)(x))[0] << 8) | ((uint8_t*)(x))[1])

bool H265PropParser::splitPropSet(const uint8_t* data, size_t length,
				std::string& vps, std::string& sps, std::string& pps)
{
	size_t start = START_CODE_LENGTH;
	size_t nextPos = findStartCode(data, length, start);
	if (nextPos == size_t(-1))
	{
		const uint8_t* p = data;
		size_t len = length;

		p += 22;
		
		int num_arrays = *p; 
		p++;

		std::string sei;
		std::string* naluData = NULL;
		for (int i = 0; i < num_arrays; i++)
		{
			int type = (*p) & 0x3f;
			p ++;

			int cnt = RB16(p);
			p += 2;

			if (type == NAL_UNIT_VPS)
			{
				naluData = &vps;
			}
			else if (type == NAL_UNIT_SPS)
			{
				naluData = &sps;
			}
			else if (type == NAL_UNIT_PPS)
			{
				naluData = &pps;
			}
			else if (type == NAL_UNIT_PREFIX_SEI)
			{
				naluData = &sei;
			}
			else if (type == NAL_UNIT_SUFFIX_SEI)
			{
				naluData = &sei;
			}
			else
			{
				return false;
			}

			for (int j = 0; j < cnt; j++)
			{
				int nalu_len = RB16(p);
				p += 2;

				naluData->append((const char*)p, nalu_len);

				p += nalu_len;
			}
		}
		return true;
	}

	vps = std::string((const char*)(data+start), nextPos - start);

	start = nextPos + START_CODE_LENGTH;
	nextPos = findStartCode(data, length, start);
	if (nextPos == size_t(-1))
	{
		return false;
	}
	sps = std::string((const char*)(data+start), nextPos - start);

	start = nextPos + START_CODE_LENGTH;
	nextPos = findStartCode(data, length, start);
	if (nextPos == size_t(-1))
	{
		nextPos = length;
	}
	pps = std::string((const char*)(data+start), nextPos - start);

	return true;
}


bool H265PropParser::parseNalu(const uint8_t* data, size_t length, NaluPacket& pkt)
{
	if ((data[0] == 0) && (data[1] == 0))
	{
		if (data[2] == 1)
		{
			pkt.prefix = 3;
		}
		else if ((data[2] == 0) && (data[3] == 1))
		{
			pkt.prefix = 4;
		}
		else
		{
			return false;
		}
	}

	pkt.data = data + pkt.prefix;
	pkt.length = length - pkt.prefix;
	pkt.type = ((pkt.data[0] >> 1) & 0x3F);

	return true;
}

bool H265PropParser::isKeyFrame(uint8_t naluType)
{
    return (naluType == NAL_IDR_W_RADL || naluType == NAL_IDR_N_LP) || (naluType >= NAL_UNIT_VPS && naluType <= NAL_UNIT_PPS);
}


} /* namespace av */
